package nl.rug.jvleusen.datastructurelib.tree;

import java.util.ArrayList;
import java.util.List;
import nl.rug.jvleusen.datastructurelib.ITree;

/**
 * This implementation of {@link ITree} implements the N-ary tree, which is a
 * tree in which each node will have at most <tt>n</tt> elements, allowing for
 * great control over the actual use of this code.
 *
 * @author Jeroen van Leusen
 */
public class NaryTree<E> implements ITree<E, NaryTree.NaryTreeNode<E>> {

    private final int k;
    private NaryTreeNode<E> root = null;

    /**
     * Creates a tree of which nodes spawned from this tree can have at most k
     * elements.
     *
     * @param k maximum amounts of elements.
     */
    public NaryTree(int k) {
        this.k = k;
    }

    /**
     * Creates a variant of this tree in which there can be unlimited amounts of
     * elements in each node.
     */
    public NaryTree() {
        this.k = -1;
    }

    /**
     * Returns how many children nodes of this tree may have, will return -1 if
     * this is unlimited.
     *
     * @return Number of children that nodes of this tree might have.
     */
    public int maxChildren() {
        return k;
    }

    /**
     * {@inheritDoc}
     */
    public NaryTreeNode<E> getRoot() {
        return root;
    }

    /**
     * Factory method to properly create nodes for a certain tree, use as
     * follows:
     * {@code NaryTreeNode<T> node = NaryTree.createNode(NaryTree<T> tree, T data)}
     *
     * @param tree Tree to build this node for.
     * @param data data to put in this node.
     * @return A constructed node for the given tree.
     */
    public static <T> NaryTreeNode<T> createNode(NaryTree<T> tree, T data) {
        return new NaryTreeNode<T>(tree, data);
    }

    /**
     * Sets the root to the given node.
     *
     * @param root new root.
     * @throws IllegalArgumentException thrown if the root is not created using
     * this tree.
     */
    public void setRoot(NaryTreeNode<E> root) {
        if (root == null || this.equals(root.tree)) {
            this.root = root;
        } else {
            throw new IllegalArgumentException("To insert into a tree, the node's tree must match the tree in question.");
        }
    }

    /**
     * Creates a list of all nodes within this tree.
     *
     * @return A list containing all nodes, will be empty if this is an empty
     * tree.
     */
    public List<NaryTreeNode<E>> toList() {
        List<NaryTreeNode<E>> list = new ArrayList<NaryTreeNode<E>>();
        if (root != null) {
            walk(root, list);
        }
        return list;
    }

    /**
     * Returns the size of this tree.
     *
     * @return the number of elements in this tree.
     */
    public int size() {
        if (root == null) {
            return 0;
        } else {
            return root.size();
        }
    }

    /**
     * {@inheritDoc}
     */
    public boolean isEmpty() {
        return root == null;
    }

    /**
     * {@inheritDoc}
     */
    public int height() {
        if (root == null) {
            return -1;
        } else {
            return root.height();
        }
    }

    private void walk(NaryTreeNode<E> element, List<NaryTreeNode<E>> list) {
        list.add(element);
        for (NaryTreeNode<E> data : element.getChildren()) {
            walk(data, list);
        }
    }

    /**
     * Implementation of {@link TreeNode} that allows for <tt>n</tt>
     * elements, based on the tree its built with.
     */
    public static class NaryTreeNode<E> implements ITree.TreeNode<E> {

        private final NaryTree<E> tree;
        private List<NaryTreeNode<E>> children;
        private E data;

        protected NaryTreeNode(NaryTree<E> tree, E data) {
            this.tree = tree;
            this.data = data;
        }

        /**
         * {@inheritDoc}
         */
        public E getData() {
            return data;
        }

        /**
         * Returns a list of all the children nodes in this node.
         *
         * @return Said list.
         */
        public List<NaryTreeNode<E>> getChildren() {
            if (children == null) {
                return new ArrayList<NaryTreeNode<E>>(0);
            }
            return this.children;
        }

        /**
         * @return How many children this node has.
         */
        public int getNumberOfChildren() {
            if (children == null) {
                return 0;
            }
            return children.size();
        }

        /**
         * Adds a node to this node as child
         *
         * @param child the node to add
         * @throws IllegalArgumentException If the node is not built with the
         * same tree as this node.
         * @throws IndexOutOfBoundsException If the node cannot have more
         * children.
         */
        public void addChild(NaryTreeNode<E> child) {
            if (!this.tree.equals(child.tree)) { //Non-matching Tree.
                throw new IllegalArgumentException("To insert into a tree, nodes need to be created using this same tree.");
            }
            if (children == null) {
                children = new ArrayList<NaryTreeNode<E>>(tree.maxChildren()); //Lazy initialization.
            }
            if (children.size() < tree.maxChildren() || tree.maxChildren() == -1) {
                children.add(child);
            } else {
                throw new IndexOutOfBoundsException(String.format(
                        "Nodes in this tree may not have more than %d children.",
                        tree.maxChildren()));
            }
        }

        /**
         * Inserts a childnode at the given index.
         *
         * @param index index to insert the node at.
         * @param child node to insert
         * @throws IllegalArgumentException If the node is not built with the
         * same tree as this node.
         * @throws IndexOutOfBoundsException If the index lies outside amount of
         * children this node may have
         */
        public void insertChildAt(int index, NaryTreeNode<E> child) {
            if (index == getNumberOfChildren()) {
                addChild(child); //Just a regular add.
            } else {
                if (!this.tree.equals(child.tree)) { //Non-matching Tree.
                    throw new IllegalArgumentException("To insert into a tree, nodes need to be created using this same tree.");
                }
                children.get(index); //just to throw the exception, and stop here
                if (index + 1 < tree.maxChildren() || tree.maxChildren() == -1) {
                    //first remove the child at that index
                    children.remove(index);
                    children.add(index, child);
                } else {
                    throw new IndexOutOfBoundsException(String.format(
                            "Nodes in this tree may not have more than %d children.",
                            tree.maxChildren()));
                }
            }
        }

        /**
         * Removes the child at the given index
         *
         * @param index index to remove.
         * @throws IndexOutOfBoundsException If thrown
         */
        public void removeChildAt(int index) {
            children.remove(index);
        }

        /**
         * Returns the child at the given index
         *
         * @param index Index to look at
         * @return The child at the given index, null if it doesn't exist.
         * @throws IndexOutOfBoundsException If thrown.
         */
        public NaryTreeNode<E> getChildAt(int index) {
            return children.get(index);
        }

        /**
         * Changes the data of this node to the given element
         *
         * @param data Element to change this node's element to.
         */
        public void setData(E data) {
            this.data = data;
        }

        /**
         * {@inheritDoc}
         */
        public boolean isLeaf() {
            return getNumberOfChildren() == 0;
        }

        /**
         * {@inheritDoc}
         */
        public int height() {
            if (this.isLeaf()) {
                return 0;
            }

            int maxHeight = 0;
            for (NaryTreeNode<E> child : children) {
                maxHeight = Math.max(maxHeight, child.height());
            }
            return maxHeight + 1;
        }

        /**
         * The number of elements within this subtree.
         *
         * @return said amount.
         */
        public int size() {
            int count = (data == null) ? 0 : 1;
            if (children != null) {
                for (NaryTreeNode<E> child : children) {
                    count += child.size();
                }
            }
            return count;
        }
    }
}
